from fastapi import APIRouter, UploadFile, File
from fastapi.responses import StreamingResponse
from ..db import session_scope, Label, TextDoc, Annotation
import json
import io

router = APIRouter(prefix="", tags=["export-import"])


@router.post("/export/project")
def export_project(payload: dict | None = None):
    """
    导出项目数据。
    - 若提供 payload（前端传入当前文档与右侧标注），仅导出该文档与这些标注；
    - 否则导出完整项目（全部 labels/docs/annotations）。
    """
    with session_scope() as s:
        labels = s.query(Label).all()

        # 优先使用前端传来的“当前文档 + 标注列表”
        if payload and (payload.get("doc") or payload.get("annotations")):
            doc = payload.get("doc") or {}
            anns = payload.get("annotations") or []
            data = {
                "labels": [
                    {
                        "id": str(l.id),
                        "parent_id": str(l.parent_id) if l.parent_id else None,
                        "name": l.name,
                        "color": l.color,
                        "description": l.description,
                        "sort": l.sort,
                    }
                    for l in labels
                ],
                "docs": [
                    {
                        "id": doc.get("id"),
                        "title": doc.get("title"),
                        "content": doc.get("content", ""),
                        "encoding": doc.get("encoding"),
                    }
                ],
                "annotations": [
                    {
                        "id": a.get("id"),
                        "doc_id": doc.get("id"),  # 可能为 None（前端未提供 id）
                        "label_id": a.get("label_id"),
                        "start_offset": a.get("start_offset"),
                        "end_offset": a.get("end_offset"),
                        "text": a.get("text", ""),
                    }
                    for a in anns
                ],
            }
            return {"code": 0, "message": "ok", "data": data}

        # 否则回退到完整导出
        docs = s.query(TextDoc).all()
        anns = s.query(Annotation).all()
        data = {
            "labels": [
                {
                    "id": str(l.id),
                    "parent_id": str(l.parent_id) if l.parent_id else None,
                    "name": l.name,
                    "color": l.color,
                    "description": l.description,
                    "sort": l.sort,
                }
                for l in labels
            ],
            "docs": [
                {
                    "id": str(d.id),
                    "title": d.title,
                    "content": d.content,
                    "encoding": d.encoding,
                }
                for d in docs
            ],
            "annotations": [
                {
                    "id": str(a.id),
                    "doc_id": str(a.doc_id),
                    "label_id": str(a.label_id),
                    "start_offset": a.start_offset,
                    "end_offset": a.end_offset,
                    "text": a.text,
                }
                for a in anns
            ],
        }
        return {"code": 0, "message": "ok", "data": data}


@router.get("/export/labels/excel")
def export_labels_excel():
    """
    将标签导出为树形 Excel（.xlsx）。
    - 动态列：每一层级三列（L{n}_Name, L{n}_Color, L{n}_Sort），便于可视化编辑层级结构；
    - 参考列：Node_ID, Parent_ID, Description 便于追踪与回导（可选用）。
    """
    # 延迟导入，避免在缺失依赖时报错到应用启动阶段
    try:
        from openpyxl import Workbook  # type: ignore
    except Exception as e:  # pragma: no cover
        return {"code": 1, "message": f"openpyxl not installed: {e}"}

    with session_scope() as s:
        # 读取并构建树
        all_labels = s.query(Label).all()
        by_parent: dict[str | None, list[Label]] = {}
        for l in all_labels:
            pid = str(l.parent_id) if l.parent_id else None
            by_parent.setdefault(pid, []).append(l)
        # 同级按 sort、name 排序
        for k in by_parent:
            by_parent[k].sort(key=lambda x: (x.sort or 0, x.name or ""))

        # DFS 计算最大层级 & 生成路径
        max_depth = 1
        paths: list[list[Label]] = []

        def dfs(node: Label, stack: list[Label]):
            nonlocal max_depth
            stack.append(node)
            max_depth = max(max_depth, len(stack))
            children = by_parent.get(str(node.id), [])
            if not children:
                paths.append(stack.copy())  # 叶子
            else:
                for ch in children:
                    dfs(ch, stack)
            stack.pop()

        roots = by_parent.get(None, [])
        for r in roots:
            dfs(r, [])

        # 若存在孤儿（parent_id 指向不存在），也单独导出作为根
        known_ids = {str(l.id) for l in all_labels}
        referenced = {str(l.parent_id) for l in all_labels if l.parent_id}
        orphans = [
            l for l in all_labels if l.parent_id and str(l.parent_id) not in known_ids
        ]
        for o in orphans:
            dfs(o, [])

        # 生成工作簿
        wb = Workbook()
        ws = wb.active
        ws.title = "labels"
        # 表头
        header: list[str] = []
        for i in range(1, max_depth + 1):
            header += [f"L{i}_Name", f"L{i}_Color", f"L{i}_Sort"]
        header += ["Node_ID", "Parent_ID", "Description"]
        ws.append(header)

        # 行：以路径为单位（每个叶子一行）；如需“每个节点都一行”，可把 paths 构造改为在 dfs 中对每个节点追加
        for path in paths:
            row: list[object | None] = []
            # 填充层级列
            for i in range(max_depth):
                if i < len(path):
                    n = path[i]
                    row += [n.name, n.color, n.sort]
                else:
                    row += [None, None, None]
            # 参考列（当前节点 = 路径最后一个）
            cur = path[-1]
            row += [
                str(cur.id),
                str(cur.parent_id) if cur.parent_id else None,
                cur.description,
            ]
            ws.append(row)

        bio = io.BytesIO()
        wb.save(bio)
        bio.seek(0)
        headers = {"Content-Disposition": 'attachment; filename="labels.xlsx"'}
        return StreamingResponse(
            bio,
            media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            headers=headers,
        )


@router.post("/export/annotations")
def export_annotations_only():
    """
    导出精简“标注内容”数据：仅包含每条标注的片段文本、所属文档标题与标签名称/颜色等，避免导出整个项目。

    返回结构：
    {
        code, message,
        data: {
            annotations: [
                {
                    id, doc_id, doc_title, label_id, label_name, label_color,
                    start_offset, end_offset, text
                }
            ]
        }
    }
    """
    with session_scope() as s:
        # 获取必要映射，避免 N+1
        labels = {str(l.id): l for l in s.query(Label).all()}
        docs = {str(d.id): d for d in s.query(TextDoc).all()}
        anns = s.query(Annotation).all()
        rows = []
        for a in anns:
            lid = str(a.label_id)
            did = str(a.doc_id)
            lobj = labels.get(lid)
            dobj = docs.get(did)
            rows.append(
                {
                    "id": str(a.id),
                    "doc_id": did,
                    "doc_title": getattr(dobj, "title", None),
                    "label_id": lid,
                    "label_name": getattr(lobj, "name", None),
                    "label_color": getattr(lobj, "color", None),
                    "start_offset": a.start_offset,
                    "end_offset": a.end_offset,
                    "text": a.text,
                }
            )
        return {"code": 0, "message": "ok", "data": {"annotations": rows}}


@router.post("/import/project")
async def import_project(file: UploadFile = File(...)):
    import uuid

    content = (await file.read()).decode("utf-8", errors="ignore")
    payload = json.loads(content)

    def parse_uuid(val):
        """解析UUID字符串，如果为None或空字符串则返回None"""
        if not val or val == "None":
            return None
        if isinstance(val, str):
            try:
                return uuid.UUID(val)
            except ValueError:
                return None
        return val

    with session_scope() as s:
        # 清空标注和文档（保留标签进行增量更新）
        s.query(Annotation).delete()
        s.query(TextDoc).delete()
        s.flush()

        # 增量导入标签：存在则更新，不存在则创建
        for l in payload.get("labels", []):
            label_id = parse_uuid(l.get("id"))
            if not label_id:
                label_id = uuid.uuid4()
            parent_id = parse_uuid(l.get("parent_id"))

            # 查找现有标签
            existing_label = s.query(Label).filter(Label.id == label_id).first()

            if existing_label:
                # 标签存在，检查是否需要更新
                need_update = False
                new_name = l.get("name", "")
                new_color = l.get("color")
                new_description = l.get("description")
                new_sort = l.get("sort", 0)

                if existing_label.name != new_name:
                    existing_label.name = new_name
                    need_update = True
                if existing_label.color != new_color:
                    existing_label.color = new_color
                    need_update = True
                if existing_label.description != new_description:
                    existing_label.description = new_description
                    need_update = True
                if existing_label.parent_id != parent_id:
                    existing_label.parent_id = parent_id
                    need_update = True
                if existing_label.sort != new_sort:
                    existing_label.sort = new_sort
                    need_update = True

                if need_update:
                    s.flush()
            else:
                # 标签不存在，创建新标签
                nl = Label(
                    id=label_id,
                    parent_id=parent_id,
                    name=l.get("name", ""),
                    color=l.get("color"),
                    description=l.get("description"),
                    sort=l.get("sort", 0),
                )
                s.add(nl)
        s.flush()

        # 导入文档
        for d in payload.get("docs", []):
            doc_id = parse_uuid(d.get("id"))
            if not doc_id:
                doc_id = uuid.uuid4()

            nd = TextDoc(
                id=doc_id,
                title=d.get("title"),
                content=d.get("content", ""),
                encoding=d.get("encoding"),
            )
            s.add(nd)
        s.flush()

        # 导入标注
        for a in payload.get("annotations", []):
            anno_id = parse_uuid(a.get("id"))
            if not anno_id:
                anno_id = uuid.uuid4()
            doc_id = parse_uuid(a.get("doc_id"))
            label_id = parse_uuid(a.get("label_id"))

            # 检查必需的外键是否存在
            if not doc_id or not label_id:
                continue  # 跳过无效的标注

            na = Annotation(
                id=anno_id,
                doc_id=doc_id,
                label_id=label_id,
                start_offset=a["start_offset"],
                end_offset=a["end_offset"],
                text=a.get("text", ""),
            )
            s.add(na)
        return {"code": 0, "message": "ok"}
